# Copyright (C) 2016 Cuckoo Foundation.
# This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org
# See the file 'docs/LICENSE' for copying permission.

from lib.cuckoo.common.abstracts import Signature

class ExploitHeapspray(Signature):
    name = "exploit_heapspray"
    description = "A potential heapspray has been detected"
    severity = 3
    categories = ["exploit"]
    authors = ["Cuckoo Technologies", "Kevin Ross"]
    minimum = "2.0"
    references = ["https://www.corelan.be/index.php/2011/12/31/exploit-writing-tutorial-part-11-heap-spraying-demystified/"]

    filter_apinames = "NtAllocateVirtualMemory",

    def init(self):
        self.mem = {}
        self.prot = {}
        self.heaptotals = dict()
        self.ignore = False
        if self.get_results("target", {}).get("category") == "file":
            f = self.get_results("target", {}).get("file", {})
            if "PE32 executable" in  f.get("type", "") or "PE32+ executable" in f.get("type", ""):
                self.ignore = True

    def on_call(self, call, process):
        pname = process["process_name"]
        protection = call["arguments"]["protection"]
        alloc_type = call["arguments"]["allocation_type"]
        region_size = call["arguments"]["region_size"]
        allocation_typefull = call["flags"].get("allocation_type")

        if "MEM_COMMIT" in allocation_typefull:
            combo = pname, region_size, protection, alloc_type
            self.mem[combo] = self.mem.get(combo, 0) + 1
            self.prot[protection] = call["flags"].get("protection")

    def on_complete(self):
        if not self.ignore:
            for combo, count in self.mem.items():
                pname, region_size, protection, alloc_type = combo

                if count >= 50:
                    written = int(region_size) * int(count)/1024/1024
                    if written >= 50:
                        if pname not in self.heaptotals:
                            self.heaptotals[pname] = 0
                        self.heaptotals[pname] += written
                        self.mark(
                            process=pname,
                            name="heapspray",
                            protection=self.prot.get(protection, protection),
                            count=count,
                            length=region_size,
                            total_mb=written,
                        )

            if self.heaptotals:
                for pname, total in self.heaptotals.items():
                    if "sprayed onto the heap of the" in self.description:
                        self.description += " and %d megabytes onto the heap of the %s process" % (total, pname)
                    else:
                        self.description += ". %d megabytes was sprayed onto the heap of the %s process" % (total, pname)
                    if total > 1024:
                        self.severity = 6
                    elif total > 512:
                        self.severity = 5
                    elif total > 256:
                        self.severity = 4

        return self.has_marks()

# Copyright (C) 2015 Optiv, Inc. (brad.spengler@optiv.com), Updated 2016 for Cuckoo 2.0
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

critical_apinames = [
    "NtAllocateVirtualMemory",
    "NtProtectVirtualMemory",
    "VirtualProtectEx",
    "NtWriteVirtualMemory",
    "NtWow64WriteVirtualMemory64",
    "WriteProcessMemory",
]

class StackPivot(Signature):
    name = "stack_pivot"
    description = "Stack pivoting was detected when using a critical API"
    severity = 6
    categories = ["exploit", "rop"]
    authors = ["Optiv", "Kevin Ross"]
    minimum = "2.0"

    filter_apinames = critical_apinames

    def __init__(self, *args, **kwargs):
        Signature.__init__(self, *args, **kwargs)
        self.ignore = False
        self.pname = []
        if self.get_results("target", {}).get("category") == "file":
            f = self.get_results("target", {}).get("file", {})
            if "PE32 executable" in f.get("type", ""):
                self.ignore = True

    def on_call(self, call, process):
        if "stack_pivoted" in call["arguments"]:
            if call["arguments"]["stack_pivoted"] == 1 and not self.ignore:
                pname = process["process_name"]
                if pname not in self.pname:
                    self.pname.append(pname)
                self.mark_call()

    def on_complete(self):
        if len(self.pname) == 1:
            self.description = "Stack pivoting was detected when using a critical API by the process "
            for pname in self.pname:
                self.description += pname
        elif len(self.pname) > 1:
            self.description = "Stack pivoting was detected when using a critical API by the processes "
            self.description += ", ".join(self.pname)
        return self.has_marks()

class DEPHeapBypass(Signature):
    name = "dep_heap_bypass"
    description = "DEP was bypassed by marking part of the heap executable following stack pivoting"
    severity = 6
    categories = ["exploit"]
    authors = ["Optiv", "Kevin Ross"]
    minimum = "2.0"

    filter_apinames = critical_apinames

    def __init__(self, *args, **kwargs):
        Signature.__init__(self, *args, **kwargs)
        self.ignore = False
        self.pname = []
        if self.get_results("target", {}).get("category") == "file":
            f = self.get_results("target", {}).get("file", {})
            if "PE32 executable" in f.get("type", ""):
                self.ignore = True

    def on_call(self, call, process):
        if "stack_dep_bypass" in call["arguments"]:
            if call["arguments"]["heap_dep_bypass"] == 1 and call["arguments"]["stack_pivoted"] == 1 and not self.ignore:
                pname = process["process_name"]
                if pname not in self.pname:
                    self.pname.append(pname)
                self.mark_call()

    def on_complete(self):
        if len(self.pname) == 1:
            self.description = "DEP was bypassed by marking part of the heap executable following stack pivoting in the process "
            for pname in self.pname:
                self.description += pname
        elif len(self.pname) > 1:
            self.description = "DEP was bypassed by marking part of the heap executable by the processes "
            self.description += ", ".join(self.pname)
        return self.has_marks()

class DEPStackBypass(Signature):
    name = "dep_stack_bypass"
    description = "DEP was bypassed by marking part of the stack executable"
    severity = 3
    categories = ["exploit"]
    authors = ["Optiv", "Kevin Ross"]
    minimum = "2.0"

    filter_apinames = critical_apinames

    def __init__(self, *args, **kwargs):
        Signature.__init__(self, *args, **kwargs)
        self.ignore = False
        self.pname = []
        if self.get_results("target", {}).get("category") == "file":
            f = self.get_results("target", {}).get("file", {})
            if "PE32 executable" in f.get("type", ""):
                self.ignore = True

    def on_call(self, call, process):
        if "stack_dep_bypass" in call["arguments"]:
            if call["arguments"]["stack_dep_bypass"] == 1 and not self.ignore:
                pname = process["process_name"]
                if pname not in self.pname:
                    self.pname.append(pname)
                self.mark_call()

    def on_complete(self):
        if len(self.pname) == 1:
            self.description = "DEP was bypassed by marking part of the stack executable by the process "
            for pname in self.pname:
                self.description += pname
        elif len(self.pname) > 1:
            self.description = "DEP was bypassed by marking part of the stack executable by the processes "
            self.description += ", ".join(self.pname)
        return self.has_marks()

class ShellcodeWriteProcessMemory(Signature):
    name = "shellcode_writeprocessmemory"
    description = "Found potential shellcode being written to a memory region previously marked executable following DEP bypass"
    severity = 3
    categories = ["exploit", "shellcode"]
    authors = ["Kevin Ross"]
    minimum = "2.0"

    filter_apinames = critical_apinames

    def __init__(self, *args, **kwargs):
        Signature.__init__(self, *args, **kwargs)
        self.ignore = False
        self.exploit = 0
        self.pname = []
        self.scpname = []
        self.memoryaddresses = []
        if self.get_results("target", {}).get("category") == "file":
            f = self.get_results("target", {}).get("file", {})
            if "PE32 executable" in f.get("type", ""):
                self.ignore = True

    def on_call(self, call, process):
        if call["api"] == "WriteProcessMemory" and self.exploit and not self.ignore:
            buf = call["arguments"]["buffer"]
            pname = process["process_name"]
            addr = call["arguments"]["base_address"]
            if pname in self.pname and addr in self.memoryaddresses and len(buf) > 0:
                self.scpname.append(pname)
                self.mark_call()

        elif "stack_pivoted" in call["arguments"]:
            if (call["arguments"]["stack_pivoted"] == 1 or call["arguments"]["stack_dep_bypass"] == 1 or call["arguments"]["heap_dep_bypass"] == 1) and not self.ignore:
                pname = process["process_name"]
                addr = call["arguments"]["base_address"]
                if pname not in self.pname:
                    self.pname.append(pname)
                if addr not in self.memoryaddresses:
                    self.memoryaddresses.append(addr)
                self.exploit = 1

    def on_complete(self):
        if len(self.scpname) == 1:
            self.description = "Found potential shellcode being written to a memory region previously marked executable following DEP bypass in the process "
            for pname in self.scpname:
                self.description += pname
        elif len(self.scpname) > 1:
            self.description = "Found potential shellcode being written to a memory region previously marked executable following DEP bypass in the processes "
            self.description += ", ".join(self.scpname)
        return self.has_marks()

class StackPivotShellcodeAPIs(Signature):
    name = "stack_pivot_shellcode_apis"
    description = "API calls following stack pivoting indicative of shellcode execution have been detected"
    severity = 3
    categories = ["exploit", "rop", "shellcode"]
    authors = ["Kevin Ross"]
    minimum = "2.0"
    evented = True
    def __init__(self, *args, **kwargs):
        Signature.__init__(self, *args, **kwargs)
        self.ignore = False
        self.pname = []
        if self.get_results("target", {}).get("category") == "file":
            f = self.get_results("target", {}).get("file", {})
            if "PE32 executable" in f.get("type", ""):
                self.ignore = True

    filter_apinames = set(["LdrLoadDll","LdrGetDllHandle","URLDownloadToFileW","URLDownloadToCacheFileW","CreateProcessInternalW","NtCreateProcess","NtCreateProcessEx","NtCreateUserProcess","RtlCreateUserProcess"])

    def on_call(self, call, process):
        if "stack_pivoted" in call["arguments"]:
            if call["arguments"]["stack_pivoted"] == 1 and not self.ignore:
                pname = process["process_name"]
                if pname not in self.pname:
                    self.pname.append(pname)
                self.mark_call()

    def on_complete(self):
        if len(self.pname) == 1:
            self.description = "API calls following stack pivoting indicative of shellcode execution have been detected in the process "
            for pname in self.pname:
                self.description += pname
        elif len(self.pname) > 1:
            self.description = "API calls following stack pivoting indicative of shellcode execution have been detected in the processes "
            list = ", ".join(self.pname )
            self.description += list
        return self.has_marks()

class StackPivotShellcodeCreateProcess(Signature):
    name = "stackpivot_shellcode_createprocess"
    description = "A process was created by shellcode following a stack pivot"
    severity = 3
    categories = ["exploit", "rop", "shellcode"]
    authors = ["Kevin Ross"]
    minimum = "2.0"
    evented = True
    def __init__(self, *args, **kwargs):
        Signature.__init__(self, *args, **kwargs)
        self.ignore = False
        self.pname = []
        if self.get_results("target", {}).get("category") == "file":
            if "PE32 executable" in self.get_results("target", {}).get("file", {}).get("type", ""):
                self.ignore = True

    filter_apinames = set(["CreateProcessInternalW","NtCreateProcess","NtCreateProcessEx","NtCreateUserProcess","RtlCreateUserProcess"])

    def on_call(self, call, process):
        if "stack_pivoted" in call["arguments"]:
            if call["arguments"]["stack_pivoted"] == 1 and not self.ignore:
                pname = process["process_name"]
                if pname not in self.pname:
                    self.pname.append(pname)
                self.mark_call()

    def on_complete(self):
        if len(self.pname) == 1:
            self.description = "A process was created by shellcode following a stack pivot by the process "
            for pname in self.pname:
                self.description += pname
        elif len(self.pname) > 1:
            self.description = "A process was created by shellcode following a stack pivot by the processes "
            list = ", ".join(self.pname )
            self.description += list
        return self.has_marks()
